package veapplib

import (
	"errors"
	"fmt"
	"golang.org/x/sys/windows/registry"
	"sync"
	"syscall"
)

type SessionDll struct {
	fnGetSessionList   *syscall.LazyProc
	fnGetProcessList   *syscall.LazyProc
	fnGetRemoteAPPList *syscall.LazyProc
	fnKillProcess      *syscall.LazyProc
	fnCreateEXE        *syscall.LazyProc
	fnSendMsg          *syscall.LazyProc
	fnLogoff           *syscall.LazyProc
	fnReturnErrInfo    *syscall.LazyProc
	fnReboot           *syscall.LazyProc

	fnSetSessionTimeLimit         *syscall.LazyProc
	fnGetSessionTimeLimit         *syscall.LazyProc
	fnEnableResetBroken           *syscall.LazyProc
	fnGetResetBroken              *syscall.LazyProc
	fnSetRemoteAppLogoffTimeLimit *syscall.LazyProc
	fnGetRemoteAppLogoffTimeLimit *syscall.LazyProc
	fnSetMaxIdleTime              *syscall.LazyProc
	fnGetMaxIdleTime              *syscall.LazyProc
	fnUpdateGP                    *syscall.LazyProc //INT UpdateGP()
	fSetSingleSessionPerUser      *syscall.LazyProc

	mu sync.Mutex
}

func NewSessionDll() *SessionDll {
	handle := syscall.NewLazyDLL("./SessionDLL.dll")
	return &SessionDll{
		fnGetSessionList:   handle.NewProc("GetSessionList"),
		fnGetProcessList:   handle.NewProc("GetProcessList"),
		fnGetRemoteAPPList: handle.NewProc("GetRemoteAPPList"),
		fnKillProcess:      handle.NewProc("KillProcess"),
		fnCreateEXE:        handle.NewProc("CreateEXE"),
		fnSendMsg:          handle.NewProc("SendMsg"),
		fnLogoff:           handle.NewProc("Logoff"),
		fnReturnErrInfo:    handle.NewProc("ReturnErrInfo"),
		fnReboot:           handle.NewProc("ReBoot"),

		fnSetSessionTimeLimit:         handle.NewProc("SetSessionTimeLimit"),
		fnGetSessionTimeLimit:         handle.NewProc("GetSessionTimeLimit"),
		fnEnableResetBroken:           handle.NewProc("EnableResetBroken"),
		fnGetResetBroken:              handle.NewProc("GetResetBroken"),
		fnSetRemoteAppLogoffTimeLimit: handle.NewProc("SetRemoteAppLogoffTimeLimit"),
		fnGetRemoteAppLogoffTimeLimit: handle.NewProc("GetRemoteAppLogoffTimeLimit"),
		fnSetMaxIdleTime:              handle.NewProc("SetMaxIdleTime"),
		fnGetMaxIdleTime:              handle.NewProc("GetMaxIdleTime"),
		fnUpdateGP:                    handle.NewProc("UpdateGP"),
		fSetSingleSessionPerUser:      handle.NewProc("SetSingleSessionPerUser"),
		mu:                            sync.Mutex{},
	}
}

func (s *SessionDll) Logoff(wsId int32) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	// 异步注销会话
	_, _, err := s.fnLogoff.Call(uintptr(wsId)) //nolint
	return err
}

func (s *SessionDll) SessionList() (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var buf = make([]byte, 65536)
	var bufsize uint32 = 0
	code, _, _ := s.fnGetSessionList.Call(newBuffer(buf), newBufferSize(&bufsize)) //nolint
	if code != 0 {
		return "", errors.New(fmt.Sprintf("获取会话列表失败：%d", code))
	}
	return string(buf[:bufsize]), nil
}

const (
	RegistryForRdpTime = "SOFTWARE\\Policies\\Microsoft\\Windows NT\\Terminal Services"
)

// SetSessionOptionsForV3 争对veapp的设置方法
func (s *SessionDll) SetSessionOptionsForV3(maxOutTime, maxSessionTime, resetBroken uint32, singleSession int) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	var (
		key registry.Key
		err error
	)
	if key, err = registry.OpenKey(registry.LOCAL_MACHINE, RegistryForRdpTime, registry.ALL_ACCESS); err != nil {
		return err
	}
	// 1. 如果是多会话，就关闭会话时间限制，设置会话断开1分钟后就注销
	if singleSession == 2 {
		// 设置多会话
		_, _, _ = s.fSetSingleSessionPerUser.Call(IntPtr(0))
		// 会话断开1分钟后销毁会话
		if err = key.SetDWordValue("MaxDisconnectionTime", 60*1000); err != nil {
			return err
		}
		// 会话时长和闲置时长都为无限制
		_, _, _ = s.fnSetSessionTimeLimit.Call(0) //nolint
		_, _, _ = s.fnSetMaxIdleTime.Call(0)      //nolint
		// 断开注销会话
		_, _, _ = s.fnEnableResetBroken.Call(IntPtr(1))
	}
	// 2. 如果是单会话，就设置会话时长和闲置时长
	if singleSession == 1 {
		// 设置单会话
		// 设置多会话
		_, _, _ = s.fSetSingleSessionPerUser.Call(IntPtr(1))
		// 如果是不自动回收的话
		if maxOutTime == maxSessionTime && maxSessionTime == 0 {
			if err = key.SetDWordValue("MaxDisconnectionTime", 0); err != nil {
				return err
			}
			_, _, _ = s.fnSetSessionTimeLimit.Call(0) //nolint
			_, _, _ = s.fnSetMaxIdleTime.Call(0)      //nolint
			_, _, _ = s.fnEnableResetBroken.Call(IntPtr(0))
		} else {
			_, _, _ = s.fnSetSessionTimeLimit.Call(uintptr(maxSessionTime)) //nolint
			_, _, _ = s.fnSetMaxIdleTime.Call(uintptr(maxOutTime))          //nolint
			// 设置注销会话
			if resetBroken == 1 {
				_, _, _ = s.fnEnableResetBroken.Call(IntPtr(1))
			} else {
				_, _, _ = s.fnEnableResetBroken.Call(IntPtr(0))
			}
			// 会话断开1分钟后销毁会话
			if err = key.SetDWordValue("MaxDisconnectionTime", 60*1000); err != nil {
				return err
			}
		}
	}
	s.fnUpdateGP.Call() //nolint
	return nil
}
